"""
质量、动量和能量方程求解器

实现两流体模型的数值求解：
1. 质量守恒方程
2. 动量守恒方程
3. 能量守恒方程
"""

import numpy as np
from .. import config as cfg
from ..models.subchannel import SubChannel

class TwoFluidSolver:
    """两流体模型求解器"""
    
    def __init__(self):
        """初始化求解器"""
        # 创建子通道网络
        self.channels = {}
        for channel_id, neighbors in cfg.CHANNEL_CONNECTIONS.items():
            self.channels[channel_id] = SubChannel(channel_id, neighbors)
            
        # 计算参数
        self.dt = cfg.TIME_STEP
        self.dz = cfg.CHANNEL_LENGTH / cfg.AXIAL_NODES
        
    def initialize_all_channels(self):
        """初始化所有子通道"""
        for channel in self.channels.values():
            channel.initialize_fields(
                cfg.INLET_PRESSURE,
                cfg.INLET_TEMP,
                cfg.MASS_FLUX
            )
            
    def solve_mass_conservation(self, channel, i):
        """
        求解质量守恒方程
        
        参数：
        channel: SubChannel, 子通道对象
        i: int, 轴向位置索引
        """
        if i == 0:  # 入口边界条件
            return
            
        # 获取流体物性
        props_l = channel.fluid_props.get_properties(
            channel.pressure[i], channel.T_l[i], 'liquid')
        props_v = channel.fluid_props.get_properties(
            channel.pressure[i], channel.T_v[i], 'vapor')
            
        if props_l is None or props_v is None:
            return
            
        # 计算轴向质量通量
        j_l = channel.alpha_l[i] * props_l['density'] * channel.u_l[i]
        j_v = channel.alpha_v[i] * props_v['density'] * channel.u_v[i]
        
        # 计算相变质量通量（简化处理）
        if channel.T_l[i] >= channel.fluid_props.get_saturation_temperature(channel.pressure[i]):
            gamma = 0.1 * j_l  # 假设10%的液相发生相变
        else:
            gamma = 0.0
            
        # 计算横向流动
        w_cross = 0.0
        for neighbor_id in channel.neighbors:
            neighbor = self.channels[neighbor_id]
            w_cross += channel.calculate_cross_flow(i, neighbor)
            
        # 更新液相分数
        j_l_prev = channel.alpha_l[i-1] * channel.u_l[i-1]
        j_l_curr = channel.alpha_l[i] * channel.u_l[i]
        channel.alpha_l[i] = channel.alpha_l[i-1] - self.dt * (
            (j_l_curr - j_l_prev) / self.dz  # 轴向对流
            + w_cross / (props_l['density'] * channel.area)  # 横向流动
            - gamma / props_l['density']  # 相变
        )
        
        # 更新气相分数
        channel.alpha_v[i] = 1.0 - channel.alpha_l[i]
        
    def solve_momentum_conservation(self, channel, i):
        """
        求解动量守恒方程
        
        参数：
        channel: SubChannel, 子通道对象
        i: int, 轴向位置索引
        """
        if i == 0:  # 入口边界条件
            return
            
        # 计算壁面摩擦力
        F_wl, F_wv = channel.calculate_wall_friction(i)
        
        # 计算界面剪切力
        F_i = channel.calculate_interfacial_force(i)
        
        # 计算压力梯度
        dp_dz = (channel.pressure[i] - channel.pressure[i-1]) / self.dz
        
        # 获取流体物性
        props_l = channel.fluid_props.get_properties(
            channel.pressure[i], channel.T_l[i], 'liquid')
        props_v = channel.fluid_props.get_properties(
            channel.pressure[i], channel.T_v[i], 'vapor')
            
        if props_l is None or props_v is None:
            return
            
        # 更新液相速度
        u_conv = (channel.u_l[i] + channel.u_l[i-1]) * 0.5  # 使用中心差分
        du_dz = (channel.u_l[i] - channel.u_l[i-1]) / self.dz
        channel.u_l[i] = channel.u_l[i-1] - self.dt * (
            u_conv * du_dz  # 对流项使用中心差分
            + (1/props_l['density']) * dp_dz  # 压力梯度
            + F_wl / (props_l['density'] * max(channel.alpha_l[i], 1e-6))  # 壁面摩擦
            + F_i / (props_l['density'] * max(channel.alpha_l[i], 1e-6))  # 界面剪切
            + 9.81  # 重力项
        )
        
        # 更新气相速度
        if channel.alpha_v[i] > 0:
            channel.u_v[i] = channel.u_v[i-1] - self.dt * (
                channel.u_v[i-1] * (channel.u_v[i] - channel.u_v[i-1]) / self.dz  # 对流项
                + (1/props_v['density']) * dp_dz  # 压力梯度
                + F_wv / (props_v['density'] * channel.alpha_v[i])  # 壁面摩擦
                - F_i / (props_v['density'] * channel.alpha_v[i])  # 界面剪切
                + 9.81  # 重力项
            )
            
    def solve_energy_conservation(self, channel, i):
        """
        求解能量守恒方程
        
        参数：
        channel: SubChannel, 子通道对象
        i: int, 轴向位置索引
        """
        if i == 0:  # 入口边界条件
            return
            
        # 计算壁面热流密度（简化处理，假设恒定热流）
        q_wall = cfg.MAX_HEAT_FLUX / (np.pi * channel.diameter)
        
        # 计算温度变化率
        dT_l_dt, dT_v_dt = channel.calculate_heat_transfer(i, q_wall)
        
        # 更新温度
        channel.T_l[i] = channel.T_l[i-1] + dT_l_dt * self.dt
        if channel.alpha_v[i] > 0:
            channel.T_v[i] = channel.T_v[i-1] + dT_v_dt * self.dt
            
        # 更新壁温（简化处理）
        channel.T_w[i] = channel.T_l[i] + q_wall / 1000  # 假设恒定传热系数
        
    def solve_one_timestep(self):
        """求解一个时间步"""
        # 遍历所有子通道
        for channel in self.channels.values():
            # 轴向推进
            for i in range(1, cfg.AXIAL_NODES):
                self.solve_mass_conservation(channel, i)
                self.solve_momentum_conservation(channel, i)
                self.solve_energy_conservation(channel, i)
                
    def run_simulation(self, n_timesteps):
        """
        运行完整的模拟计算
        
        参数：
        n_timesteps: int, 时间步数
        """
        print("开始数值计算...")
        
        # 初始化
        self.initialize_all_channels()
        
        # 时间推进
        for n in range(n_timesteps):
            if n % 100 == 0:
                print(f"计算第{n}个时间步...")
            self.solve_one_timestep()
            
        print("计算完成") 